# Created by Gao Song on 2018-05-02

import pandas as pd
from datetime import datetime, date, timedelta
from .rq import rq


def get_trading_dates(start_date, end_date):
    return rq.get_trading_dates(start_date, end_date)


def normalize_datetime(d):
    if isinstance(d, datetime):
        d = d.replace(hour=0, minute=0, second=0, microsecond=0)
    elif isinstance(d, date):
        d = datetime.combine(d, datetime.min.time())
    elif isinstance(d, pd.Timestamp):
        d = d.to_pydatetime().replace(hour=0, minute=0, second=0, microsecond=0)
    return d

def get_prev_trading_date(d):
    return normalize_datetime(rq.get_previous_trading_date(d, n=1))

def get_next_trading_date(d):
    return normalize_datetime(rq.get_next_trading_date(d, n=1))

def is_trading_day(d):
    d = normalize_datetime(d)
    return len(get_trading_dates(d, d)) == 1


def holiday_eve(d):
    d = normalize_datetime(d)
    if is_trading_day(d):
        if len(rq.get_trading_dates(start_date=d+timedelta(days=1), end_date=d+timedelta(days=3))) == 0:
            return True
        if (not is_trading_day(d+timedelta(days=1))) and (d + timedelta(days=1)).isoweekday() in [1, 2, 3, 4, 5]:
            return True
    else:
        return False


if __name__ == '__main__':
    from datetime import datetime, date

    assert holiday_eve(datetime(2019, 12, 31))
    assert not holiday_eve(datetime(2020, 1, 1))
    assert holiday_eve(datetime(2020, 1, 23))
    assert holiday_eve(datetime(2019, 9, 12))
    assert not holiday_eve(datetime(2020, 1, 17))
    assert not holiday_eve(datetime(2020, 1, 24))

    assert holiday_eve(date(2019, 12, 31))
    assert not holiday_eve(date(2020, 1, 1))
    assert holiday_eve(date(2020, 1, 23))
    assert holiday_eve(date(2019, 9, 12))
    assert not holiday_eve(date(2020, 1, 17))
    assert not holiday_eve(date(2020, 1, 24))

    assert get_prev_trading_date(datetime(2021, 5, 11)) == datetime(2021, 5, 10)
    assert get_prev_trading_date(datetime(2021, 5, 11, 20, 1, 2)) == datetime(2021, 5, 10)
    assert get_prev_trading_date(datetime(2021, 5, 6)) == datetime(2021, 4, 30)
    assert get_prev_trading_date(datetime(2021, 5, 5)) == datetime(2021, 4, 30)
